import os
import re
from collections import defaultdict
from itertools import combinations

# 数据预处理：读取数据并转换为事务格式
def load_transactions(data_dir):
    transactions = []
    for category_dir in os.listdir(data_dir):
        category_path = os.path.join(data_dir, category_dir)
        if os.path.isdir(category_path):
            for file_name in os.listdir(category_path):
                file_path = os.path.join(category_path, file_name)
                with open(file_path, 'r', encoding='latin1') as f:
                    # 提取文章内容，去掉标点并分词
                    content = f.read().lower()
                    words = re.findall(r'\b\w+\b', content)
                    transactions.append(set(words))
    return transactions

# Apriori算法实现
def apriori(transactions, min_support):
    def get_frequent_itemsets(candidates, transactions, min_support):
        item_counts = defaultdict(int)
        for transaction in transactions:
            for candidate in candidates:
                if candidate.issubset(transaction):
                    item_counts[candidate] += 1
        return {
            item: count
            for item, count in item_counts.items()
            if count / len(transactions) >= min_support
        }
    
    # 初始单项集
    items = {frozenset([item]) for transaction in transactions for item in transaction}
    transactions_count = len(transactions)
    
    # 生成频繁项集
    frequent_itemsets = []
    k = 1
    while items:
        print(f"Generating {k}-itemsets...")
        frequent_items = get_frequent_itemsets(items, transactions, min_support)
        if not frequent_items:
            break
        frequent_itemsets.append(frequent_items)
        # 生成候选项集
        items = {item1.union(item2) for item1 in frequent_items for item2 in frequent_items if len(item1.union(item2)) == k + 1}
        k += 1

    return frequent_itemsets

# 主函数
def main():
    # 数据集路径
    data_dir = "./20newsgroups_data/20news-18828"
    
    # 加载事务数据
    print("Loading transactions...")
    transactions = load_transactions(data_dir)
    print(f"Loaded {len(transactions)} transactions.")
    
    # Apriori算法
    print("Running Apriori algorithm...")
    min_support = 0.05  # 最小支持度阈值
    frequent_itemsets = apriori(transactions, min_support)
    
    # 输出结果
    for k, itemsets in enumerate(frequent_itemsets, start=1):
        print(f"Frequent {k}-itemsets:")
        for itemset, count in itemsets.items():
            print(f"  {set(itemset)}: {count}")
    
if __name__ == "__main__":
    main()
